'''
Created on Nov 30, 2010

@author: oabalbin
'''

import exome.gatk_cluster.gatk_realignment_cluster as gra
import exome.gatk_cluster.picard_commands_cluster as mypicard
import exome.gatk_cluster.samtools_commands_cluster as mysam
import exome.gatk_cluster.base_quality_calibration_cluster as gqc
import exome.gatk_cluster.cluster_jobs_header as jh

from collections import defaultdict, deque


def lane_level_realignment(sample_lanes_dict, gatk_run_dict):
    '''
    Input: a sample lane dictionary with key:sample, values: lanes that belong to that sample
    By default this realignment method uses only known indel Sites to realign. 
    The known indel sites are provided in the indeldb_file
    gatk_run_dict is a dictionary with the parameters to run the gatk realignment walker
        indeldb_file
        
    
    '''
    # gatk parameters
    '''
    ref_genome='/exds/projects/alignment_indexes/gatk/hg19/hg19.fa'
    ref_dbsnp='/exds/projects/alignment_indexes/gatk/hg19/dbsnp132_00-All_processed5.processed.vcf'
    bam_file_name='/exds/users/oabalbin/projects/snps/exomes/aM18/test/s_3_12_sequence.hg19.aln2.rmdup.sorted.bam'
    use_mem=24
    num_cores=6
    recal_analysis_outputdir='/exds/users/oabalbin/projects/snps/exomes/aM18/test/recal_analysis/'
     
    gatk_run_dict = {'path_to_gatk':'/exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/', 
                     'resources_folder':'/exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/resources/', 
                     'rscipt_path':'/exds/sw/local/R-project/bin/Rscript'}

    '''
    
    
    # Run parameters
    use_mem, num_cores = gatk_run_dict['use_mem'], gatk_run_dict['num_cores']
    path_to_gatk, path_to_picard = gatk_run_dict['path_to_gatk'], gatk_run_dict['path_to_picard']
    ref_genome, snpdb_file, indeldb_file = gatk_run_dict['ref_genome'], gatk_run_dict['snpdb_file'], gatk_run_dict['indeldb_file'] 
    path_to_intervals = gatk_run_dict['path_to_intervals']
    recal_analysis_outputdir=gatk_run_dict['recal_analysis_outputdir']
    temp_dir, out_dir = gatk_run_dict['temp_dir'],gatk_run_dict['out_dir']
    
    # cluster parameters
    new_sample_lanes_dict=defaultdict(deque)
    allwt,qcwt= str('1:00:00'), str('20:00')
    # For not multi-thread jobs
    snode, sncore, nnodes, maxcores = 1,1,1,12
 
    First=True
    for sp, thlanes in sample_lanes_dict.iteritems():    
        # Cluster parameters
        jname='rp_'+str(sp)
        nnodes, ncores = len(thlanes), sncore
        qsubfile_name=gatk_run_dict['qsubfile']+sp+'.qsub'
        qsubfile=open(qsubfile_name,'w')


        if First:
            jh.write_pbs_header(qsubfile, jname, snode, sncore, allwt)
            qsubfile.write('# Create realignment intervals at known indel positions\n')
            # gatk does not support multi-core here yet
            command, intervals_to_realign = gra.intvTorealign_only_knownSites(ref_genome, use_mem, num_cores, 
                                                                 path_to_intervals, path_to_gatk, 
                                                                 indeldb_file)
            
            qsubfile.write(command+'\n')
            First=False
                    

        for indexed_bam_file in thlanes:
            
            qsubfile.write('# Realign Bam files at indel sites. Usually using indel knownSites db\n')
            # gatk does not support multi-core here yet
            command, realigned_bam_file = gra.realign_bam_only_knownSites(ref_genome, indexed_bam_file, 
                                                           intervals_to_realign, indeldb_file, 
                                                           use_mem, num_cores, temp_dir,path_to_gatk)
            
            qsubfile.write(command+'\n')
            
            qsubfile.write('# Fix Mate information in the Realigned Bam file\n')
            command, matefixed_bam = mypicard.fixmateInformation(realigned_bam_file, use_mem, path_to_picard, temp_dir)
            qsubfile.write(command+'\n')
            
            qsubfile.write('# Mark duplicates in the Realigned Bam file\n')
            command, dedup_bam = mypicard.markDuplicates(matefixed_bam, use_mem, path_to_picard)
            qsubfile.write(command+'\n')
            
            ##########
            qsubfile_qcname=gatk_run_dict['qsubfile']+sp+'_qcstep.qsub'
            qsubfile.write('# Recalibrate the realigned Bam file\n')
            qsubfile.write('qsub '+str(qsubfile_qcname)+'\n')
            
            # Create a new pbs script for running this job            
            qc_command_list, recal_bam_file = gqc.main_baseQ_recalibration(ref_genome, snpdb_file, 
                                                               dedup_bam, use_mem, num_cores,
                                                               gatk_run_dict, recal_analysis_outputdir)
            # Create a new pbs script for running this job
            qsubfile_qc=open(qsubfile_qcname,'w')
            jh.write_pbs_header(qsubfile_qc, jname+'qc', maxcores, nodecores, qcwt)
            
            [qsubfile_qc.write(cline+'\n') for cline in qc_command_list]
            qsubfile_qc.close()
            
            # new dictionary with the realigned and 
            # recalibrated bam files paths for each lane that belongs to sample
            new_sample_lanes_dict[sp].append(recal_bam_file)
        
        
        recalibrated_lanes = new_sample_lanes_dict[sp]
        sample_recal_merged_bam = out_dir+sp+'_merged_recal.bam'
        qsubfile.write('# Merge realigned, recalibrated lanes that belong to sample'+sp+'\n')
        command = mysam.merge_bam_files(list(recalibrated_lanes), sample_recal_merged_bam)        
        qsubfile.write(command+'\n')
        qsubfile.write('# Mark duplicates in the merged Bam file\n')
        command, merged_dedup_bam = mypicard.markDuplicates(sample_recal_merged_bam, use_mem, path_to_picard)
        qsubfile.write(command+'\n')
           
    qsubfile.close()
    
    return  merged_dedup_bam

  

sample_lanes_dict={'aM18':['/exds/users/oabalbin/projects/snps/exomes/aM18/test/s_3_12_sequence.hg19.aln2.rmdup.sorted.bam']}

gatk_run_dict = {'path_to_gatk':'/exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/',
                 'path_to_picard':'/exds/sw/bioinfo/alignment/picard/picard-tools-1.35/', 
                 'resources_folder':'/exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/resources/', 
                 'rscipt_path':'/exds/sw/local/R-project/bin/Rscript',
                 'use_mem':8, 'num_cores':8,
                 'ref_genome':'/exds/projects/alignment_indexes/gatk/hg19/hg19.fa', 
                 'snpdb_file':'/exds/projects/alignment_indexes/gatk/hg19/dbsnp132_00-All_processed5.processed.vcf',
                 'indeldb_file':'/exds/projects/alignment_indexes/gatk/hg19/dbsnp132_00-All_processed5.processed.vcf',
                 'path_to_intervals':'/exds/users/oabalbin/projects/snps/exomes/aM18/test/',
                 'recal_analysis_outputdir':'/exds/users/oabalbin/projects/snps/exomes/aM18/test/recal_analysis/',
                 'temp_dir':'/exds/users/oabalbin/projects/snps/exomes/aM18/test/temp/',
                 'qsubfile':'/exds/users/oabalbin/projects/snps/exomes/aM18/test/',
                 'out_dir':'/exds/users/oabalbin/projects/snps/exomes/aM18/test/'
                 }




merged_dedup_bam = lane_level_realignment(sample_lanes_dict,gatk_run_dict)
print merged_dedup_bam

'''
for each lane.bam
    realigned.bam <- realign(lane.bam) [at only known sites]
    dedup.bam <- MarkDuplicate(realigned.bam)
    recal.bam <- recal(dedup.bam)

for each sample
    recals.bam <- merged lane-level recal.bam's for sample
    dedup.bam <- MarkDuplicates(recals.bam)


for each lane.bam
    realigned.bam <- realign(lane.bam) [at only known sites]
    dedup.bam <- MarkDuplicate(realigned.bam)
    recal.bam <- recal(dedup.bam)

for each sample
    recals.bam <- merged lane-level recal.bam's for sample
    dedup.bam <- MarkDuplicates(recals.bam)
    realigned.bam <- realign(dedup.bam) [with known sites if possible]

'''
